import numpy as np
from dataclasses import dataclass
from typing import List, Tuple
import time
import random

@dataclass
class Position:
    """3D position in the environment"""
    x: float
    y: float
    z: float
    
    def distance_to(self, other: 'Position') -> float:
        return np.sqrt(
            (self.x - other.x)**2 + 
            (self.y - other.y)**2 + 
            (self.z - other.z)**2
        )

@dataclass
class Object3D:
    """Physical object in 3D space"""
    position: Position
    size: Tuple[float, float, float]  # width, height, depth
    velocity: Tuple[float, float, float]  # vx, vy, vz
    type: str
    properties: dict

class World:
    """3D environment with physical properties"""
    def __init__(self, size=(10, 10, 5)):
        self.size = size
        self.has_gravity = True
        self.terrain = {}
        self.platforms = []
        self.obstacles = []
        self.visibility = 1.0
        self.temperature = 22
        self.humidity = 0.5
        
        # Initialize terrain
        self.terrain = self.generate_terrain()
        
        # Basic properties
        self.provider = "Suanfamama"
        self.sub_class = "3D_env"
        self.geo_location_region = "north america"
        
        # Space-time properties
        self.time = time.time()  # Current time
        self.time_step = 0.1     # Time step for updates
        self.gravity = -9.81     # Gravity in m/s^2
        self.boundaries = {
            'x': (-100, 100),    # x boundaries
            'y': (-100, 100),    # y boundaries
            'z': (0, 100)        # z boundaries (ground at z=0)
        }
        
        # Environment state
        self.objects: List[Object3D] = []
        self.light_level = 1.0   # 0 to 1
        self.air_pressure = 1.0  # atm
        
        # Agent-related properties
        self.energy = 100
        self.knowledge = 0
        self.intelligence = 0
        self.complex_level = 0
        
        # Security and privacy
        self.security_provider = ""
        self.security = 0
        self.privacy_provider = ""
        self.privacy = 0
        
    def generate_terrain(self):
        """Generate terrain height map"""
        terrain = {}
        for x in range(self.size[0]):
            for y in range(self.size[1]):
                terrain[(x, y)] = random.uniform(0, 1)
        return terrain

    def _point_in_box(self, point, box_pos, box_size):
        """Check if a point is inside a box"""
        x, y, z = point
        bx, by, bz = box_pos
        w, h, d = box_size
        
        return (bx <= x < bx + w and
                by <= y < by + h and
                bz <= z < bz + d)

    def is_valid_position(self, position):
        """Check if position is valid"""
        x, y, z = position
        
        # Check boundaries
        if not (0 <= x < self.size[0] and 
                0 <= y < self.size[1] and 
                0 <= z < self.size[2]):
            return False
        
        # Check collisions with obstacles
        if self.check_collision(position):
            return False
        
        # Check if position is supported (on ground or platform)
        if z == 0:  # Ground level is always valid
            return True
        
        # Check if on any platform
        for platform in self.platforms:
            if (platform['x'] <= x < platform['x'] + platform['width'] and
                platform['y'] <= y < platform['y'] + platform['length']):
                if z == platform['z']:  # Must be exactly at platform height
                    return True
                elif z > platform['z']:  # Above platform is invalid
                    return False
                
        return False  # No support found

    def check_collision(self, position):
        """Check for collisions with obstacles"""
        for obstacle in self.obstacles:
            if self._point_in_box(position, obstacle['position'], obstacle['size']):
                return True
        return False

    def add_object(self, obj: Object3D):
        """Add a new object to the environment"""
        self.objects.append(obj)
        
    def remove_object(self, obj: Object3D):
        """Remove an object from the environment"""
        if obj in self.objects:
            self.objects.remove(obj)
            
    def update(self, dt: float = None):
        """Update environment state based on time step"""
        if dt is None:
            dt = self.time_step
            
        current_time = time.time()
        elapsed = current_time - self.time
        self.time = current_time
        
        # Update object positions based on physics
        for obj in self.objects:
            # Update position based on velocity
            obj.position.x += obj.velocity[0] * elapsed
            obj.position.y += obj.velocity[1] * elapsed
            obj.position.z += obj.velocity[2] * elapsed + 0.5 * self.gravity * elapsed**2
            
            # Apply gravity to vertical velocity
            vx, vy, vz = obj.velocity
            obj.velocity = (vx, vy, vz + self.gravity * elapsed)
            
            # Basic collision with ground
            if obj.position.z < 0:
                obj.position.z = 0
                obj.velocity = (vx, vy, abs(vz) * 0.8)  # Bounce with energy loss
                
            # Keep objects within boundaries
            obj.position.x = np.clip(obj.position.x, *self.boundaries['x'])
            obj.position.y = np.clip(obj.position.y, *self.boundaries['y'])
            obj.position.z = np.clip(obj.position.z, *self.boundaries['z'])
        
        # Update environmental conditions based on time
        self.light_level = 0.5 + 0.5 * np.sin(2 * np.pi * (self.time % 86400) / 86400)  # Day/night cycle
        self.temperature += np.random.normal(0, 0.1)  # Small random temperature changes
        
    def get_objects_in_radius(self, center: Position, radius: float) -> List[Object3D]:
        """Get all objects within a certain radius of a point"""
        return [
            obj for obj in self.objects 
            if obj.position.distance_to(center) <= radius
        ]
        
    def reward(self, magnitude: float = 1.0):
        """Update environment state based on reward"""
        max_value = 100
        self.intelligence = min(self.intelligence + magnitude, max_value)
        self.complex_level = min(self.complex_level + magnitude, max_value)
        self.knowledge = min(self.knowledge + magnitude, max_value)
        self.energy = min(self.energy + magnitude, max_value)
        self.security = min(self.security + magnitude, max_value)
        self.privacy = min(self.privacy + magnitude, max_value)

    def create_basic_environment(self):
        """Create a basic 3D environment with some objects"""
        # Add ground plane
        ground = Object3D(
            position=Position(0, 0, 0),
            size=(200, 200, 0.1),
            velocity=(0, 0, 0),
            type="ground",
            properties={"material": "grass"}
        )
        self.add_object(ground)
        
        # Add some random objects
        for _ in range(10):
            obj = Object3D(
                position=Position(
                    x=np.random.uniform(*self.boundaries['x']),
                    y=np.random.uniform(*self.boundaries['y']),
                    z=np.random.uniform(0, 50)
                ),
                size=(1, 1, 1),
                velocity=(
                    np.random.uniform(-1, 1),
                    np.random.uniform(-1, 1),
                    np.random.uniform(0, 2)
                ),
                type="dynamic_object",
                properties={"mass": 1.0}
            )
            self.add_object(obj)

    def add_platform(self, platform):
        """Add a platform to the environment"""
        if self._is_valid_platform_config(platform):
            self.platforms.append(platform)
        else:
            raise ValueError("Invalid platform configuration")

    def add_obstacle(self, obstacle):
        """Add an obstacle to the environment"""
        if self._is_valid_obstacle_config(obstacle):
            self.obstacles.append(obstacle)
        else:
            raise ValueError("Invalid obstacle configuration")

    def set_lighting(self, condition):
        """Set lighting condition and update visibility"""
        lighting_configs = {
            'natural': 1.0,
            'dim': 0.6,
            'dark': 0.3
        }
        self.visibility = lighting_configs.get(condition, 1.0)

    def apply_gravity(self, position):
        """Apply gravity effect to a position"""
        x, y, z = position
        
        # Find highest supporting surface (platform or ground)
        max_height = 0
        
        # Check platforms
        for platform in self.platforms:
            if (platform['x'] <= x < platform['x'] + platform['width'] and
                platform['y'] <= y < platform['y'] + platform['length']):
                if platform['z'] < z and platform['z'] > max_height:
                    max_height = platform['z']
        
        # Return position with gravity applied
        return (x, y, max_height)

    def _is_valid_platform_config(self, platform):
        """Validate platform configuration"""
        required_keys = {'x', 'y', 'z', 'width', 'length'}
        if not all(key in platform for key in required_keys):
            return False
            
        # Check boundaries
        x, y, z = platform['x'], platform['y'], platform['z']
        width, length = platform['width'], platform['length']
        
        return (0 <= x + width <= self.size[0] and 
                0 <= y + length <= self.size[1] and 
                0 <= z < self.size[2])

    def _is_valid_obstacle_config(self, obstacle):
        """Validate obstacle configuration"""
        if 'position' not in obstacle or 'size' not in obstacle:
            return False
            
        x, y, z = obstacle['position']
        w, h, d = obstacle['size']
        
        return (0 <= x + w <= self.size[0] and 
                0 <= y + h <= self.size[1] and 
                0 <= z + d <= self.size[2])

    def constrain_position(self, position):
        """Constrain position within world boundaries"""
        x, y, z = position
        return (
            min(max(x, 0), self.size[0] - 1),
            min(max(y, 0), self.size[1] - 1),
            min(max(z, 0), self.size[2] - 1)
        ) 